import time
import random
from create_sub_file.write_pbs.write_pbs_common import *

def get_tiemstr():
    timestr = time.strftime('%Y%m%dT%H%M%S',time.localtime())
    return timestr

def get_free_nodes(nodeinfo):
    if "free_nodes" not in nodeinfo:
        print("Free nodes created erro!")
        exit(0)
    else:
        return nodeinfo["free_nodes"]


def get_nodes_info(process, subinfo):
    nodeinfo = subinfo["nodeinfo"]
    node_list_order = nodeinfo["node_list_order"]
    core_per_node = subinfo["cores_per_node"]
    
    # 计算node个数
    numnodes = process // core_per_node
    if process % core_per_node != 0:
        numnodes += 1
    
    if node_list_order == "normal":
        nodes = numnodes
    elif node_list_order == "random":
        nodes = []
        free_nodes = get_free_nodes(nodeinfo)
        node_index = list(range(len(free_nodes)))
        node_index = random.sample(node_index, len(node_index))  
        for i in node_index[0:numnodes]:
            nodes.append(free_nodes[i])
    elif node_list_order == "next":
        nodes = []
        free_nodes = get_free_nodes(nodeinfo)
        nodes = free_nodes[0:numnodes]
    elif node_list_order == "far":
        nodes = []
        free_nodes = get_free_nodes(nodeinfo)
        interval = len(free_nodes) / numnodes
        nodes = free_nodes[0:numnodes:interval]
    return nodes

def create_process_list(processinfo):
    process_list = []
    np_start = processinfo["np_start"]
    np_end   = processinfo["np_end"]
    interval = 0
    common_ratio = 0
    if processinfo["interval"]:
        interval = processinfo["interval"]
    elif processinfo["common_ratio"]:
        common_ratio = processinfo["common_ratio"]
    else:
        print("Please select one type to create process list.")
        exit(0)
        
    if(np_end == np_start):
        process_list = [np_start]
    elif(interval):
        if(np_start == 1):
            process_list = list(range(np_start - 1, np_end + 1, interval))
            process_list[0] = 1
        else:
            process_list = list(range(np_start, np_end + 1, interval))
    elif(common_ratio):
        while(np_start < np_end):
            process_list.append(np_start)
            np_start = np_start * common_ratio
    if(process_list[-1] != np_end):
        process_list.append(np_end)
    process_list.sort(reverse = True)

    return process_list

def tmp2run_jxpamg_dict(tmp_amg_dict):
    jxpamg_dict = {}
    except_list = get_except_list()
    for i in tmp_amg_dict:
        if i not in except_list:
            jxpamg_dict[i] = tmp_amg_dict[i]
    return jxpamg_dict

def create_mpi_option(tmp_amg_dict,runenvironments):
    mpi_dict = runenvironments["mpi"]
    numprocess = tmp_amg_dict["numprocess"]
    if mpi_dict["mpirun_option"] is not None:
        mpi_option = "mpirun {0}".format(mpi_dict["mpirun_option"])
    else:
        mpi_option = "mpirun"
    mpi_option += f" -np {numprocess}"
    return mpi_option

def create_hpctoolkit_option(runenvironments):
    if "use_hpctoolkit" in runenvironments:
        hpctoolkit_option = "hpcrun "
        hpctoolkit_env = runenvironments["hpctoolkit"]["event"]
        for i in hpctoolkit_env:
            hpctoolkit_option += f"-e {i} "
        hpctoolkit_option.strip(" ")
        if runenvironments["hpctoolkit"]["data_dir_name"] is not None:
            hpctoolkit_option += "-o ./{0}".format(runenvironments["hpctoolkit"]["data_dir_name"])
    else:
        hpctoolkit_option = ""
    return hpctoolkit_option

def create_other_para(run_parameter):
    other_para = ""
    if run_parameter["other_runpara"] is not None:
            other_para += run_parameter["other_runpara"]
    return other_para

def change_hypre_dict(rawpara, corpara, hypre_dict):
    new_hypre_dict = {}
    corpara_val_list = list(corpara.values())
    corpara_name_list = list(corpara)

    j = 0
    for i in hypre_dict:
        index = rawpara[i].index(hypre_dict[i])
        name = corpara_name_list[j]
        val  = corpara_val_list[j][index]
        if i == 'ct':
            new_hypre_dict[val] = ""
        else:
            new_hypre_dict[name] = val
        j += 1
    
    # print(new_hypre_dict)
    return new_hypre_dict

def tmp2run_hypre_dict(tmp_amg_dict,amgpara):
    hypre_dict = {}
    except_list = get_except_list()

    rawpara = amgpara["run_parameter"]
    corpara = amgpara["HYPRE_corresponding_parameter"]

    for i in tmp_amg_dict:
        if i not in except_list:
            hypre_dict[i] = tmp_amg_dict[i]

    hypre_dict = change_hypre_dict(rawpara, corpara, hypre_dict)

    return hypre_dict